'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 5_llm_self_repair.py
* @Time: 2025/10/29
* @All Rights Reserve By Brtc
'''
from typing import Any

import dotenv
from langchain_core.messages import ToolCall, AIMessage, ToolMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()


class CustomToolException(Exception):
    """自定义的工具错误异常"""

    def __init__(self, tool_call: ToolCall, exception: Exception) -> None:
        super().__init__()
        self.tool_call = tool_call
        self.exception = exception


@tool
def complex_tool(int_arg: int, float_arg: float, dict_args:dict) -> int:
    """使用复杂工具进行复杂计算操作"""
    print(dict_args)
    return int_arg * float_arg


def tool_custom_exception(msg: AIMessage, config: RunnableConfig) -> Any:
    print(msg.tool_calls[0]["args"])
    try:

        return complex_tool.invoke(msg.tool_calls[0]["args"], config)
    except Exception as e:
        raise CustomToolException(msg.tool_calls[0], e)


def exception_to_messages(inputs: dict) -> dict:
    # 1.从输入中提取错误信息
    print(inputs)
    exception = inputs.pop("exception")
    # 2.将历史消息添加到原始输入中，以便模型直到它在上一次工具调用中犯了一个错误
    print(exception.tool_call["id"])
    messages = [
        AIMessage(content="", tool_calls=[exception.tool_call]),
        ToolMessage(tool_call_id=exception.tool_call["id"], content=str(exception.exception)),
        HumanMessage(content="最后一次工具调用引发了异常，调用参数应该是5，2.1, {'a':'c'}")
    ]
    inputs["last_output"] = messages
    return inputs


# 1.创建prompt，并预留占位符，用于存储错误输出信息
prompt = ChatPromptTemplate.from_messages([
    ("human", "{query}"),
    ("placeholder", "{last_output}"),
])

# 2.创建大语言模型并绑定工具
llm = ChatOpenAI(model="gpt-4o", temperature=0).bind_tools(tools=[complex_tool], tool_choice="any")

# 3.创建链并执行工具
chain = prompt | llm | tool_custom_exception
self_correcting_chain = chain.with_fallbacks(
    [exception_to_messages | chain], exception_key="exception",
)

# 4.调用自我纠正链完成任务
print(self_correcting_chain.invoke({"query": "使用复杂工具，对应参数为5和2.1"}))